# data/mysql_db.py
"""
SQLAlchemy 的核心优势

分层架构：提供Core和ORM两个独立但相互配合的层次
数据库无关性：支持几乎所有主流关系型数据库
灵活的模型设计：灵活定义表关系、继承等高级特性
Session管理：强大的事务和对象管理
表达式语言：Python风格的SQL构建系统
性能优化：延迟加载、即时查询编译等优化策略
类型系统：丰富的类型处理和验证
"""

from datetime import datetime

import pandas as pd
from sqlalchemy import create_engine, Column, String, Integer, Date, Numeric, Text, BigInteger, DateTime
from sqlalchemy.orm import sessionmaker, declarative_base

from config.config import config

# 创建基类
Base = declarative_base()


class StockBasicInfo(Base):
    """
    股票基本信息
    """
    __tablename__ = 'stock_basic_info'

    id = Column(Integer, primary_key=True)
    # TS代码
    ts_code = Column(String(20))
    # 股票代码
    symbol = Column(String(20), unique=True, nullable=False)
    # 股票名称
    name = Column(String(20))
    # 地域
    area = Column(String(20))
    # 所属行业
    industry = Column(String(20))
    # 拼音缩写
    cnspell = Column(String(20))
    # 市场类型（主板 / 创业板 / 科创板 / CDR）
    market = Column(String(20))
    # 上市日期
    list_date = Column(Date)

    __table_args__ = {'mysql_engine': 'InnoDB', 'mysql_charset': 'utf8mb4'}

    def __repr__(self):
        return '<StockBasicInfo %r>' % self.symbol


class StockDailyData(Base):
    __tablename__ = 'stock_daily_data'

    id = Column(Integer, primary_key=True)
    # 股票代码
    ts_code = Column(String(20))
    # stock_id = Column(String(20))
    # 交易日期
    trade_date = Column(Date)
    # 开盘价
    open = Column(Numeric(10, 2))
    # 最高价
    high = Column(Numeric(10, 2))
    # 最低价
    low = Column(Numeric(10, 2))
    # 收盘价
    close = Column(Numeric(10, 2))
    # 成交量（手）
    vol = Column(Integer)
    # 成交额（千元）
    amount = Column(Numeric(18, 2))

    __table_args__ = {'mysql_engine': 'InnoDB', 'mysql_charset': 'utf8mb4'}


class Post(Base):
    __tablename__ = 'post'

    id = Column(Integer, primary_key=True)
    # 评论Id
    post_id = Column(BigInteger, unique=True, nullable=False)
    forum_id = Column(Integer, nullable=False)
    content = Column(Text, nullable=False)
    image = Column(Text, nullable=False)
    publish_time = Column(DateTime, nullable=False)

    __table_args__ = {'mysql_engine': 'InnoDB', 'mysql_charset': 'utf8mb4'}

    def __repr__(self):
        return f"""
id: {self.post_id}
内容: {self.content}
图片: {self.image}
发布时间: {self.publish_time}
{'-' * 50}
        """


class DataManager:
    def __init__(self):
        db_config = config.get('database', {}).get('mysql', {})
        db_url = f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['db_name']}?charset=utf8mb4"
        # 创建引擎
        # Engine: 管理数据库连接池并处理与数据库的通信
        self.engine = create_engine(db_url,  # dialect+driver://username:password@host:port/database
                                    echo=True,  # 启用SQL日志
                                    pool_size=db_config['pool_size'],  # 连接池大小
                                    max_overflow=10,  # 连接池最大溢出
                                    pool_timeout=30,  # 连接超时时间
                                    pool_recycle=3600  # 连接回收时间
                                    )
        # MetaData: 包含数据库模式信息的集合，记录表和列的定义。
        # 创建表
        Base.metadata.create_all(self.engine)
        # 创建会话工厂
        # Session: 管理对象的持久化和事务。
        self.Session = sessionmaker(bind=self.engine)

    def get_session(self):
        return self.Session()

    def get_stock_daily_data(self, ts_code: str) -> pd.DataFrame:
        """获取股票数据并设置索引为交易日期"""
        session = self.get_session()
        try:
            # 使用 ORM 查询替代原始 SQL 字符串
            query = (session.query(StockDailyData).filter_by(ts_code=ts_code)
                     .where(StockDailyData.trade_date.between(datetime(2025, 1, 1), datetime.now()))
                     .order_by(StockDailyData.trade_date))
            # 转换为 DataFrame
            df = pd.read_sql(query.statement, query.session.bind)

            # 设置索引并处理缺失值
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df.set_index('trade_date', drop=False, inplace=True)
            df.ffill(inplace=True)

            return df

        except Exception as e:
            print(f"获取股票数据失败: {str(e)}")
            return pd.DataFrame()

        finally:
            session.close()

    def save_stock_daily_data(self, df: pd.DataFrame):
        """
        将DataFrame数据批量存储到MySQL

        参数:
            df (DataFrame): 需要存储的数据
        """
        if df.empty:
            return

        session = self.get_session()
        try:
            # 转换字段名和类型
            records = []
            for _, row in df.iterrows():
                records.append(
                    StockDailyData(
                        ts_code=str(row['ts_code']),
                        trade_date=datetime.strptime(str(row['trade_date']), '%Y%m%d').date(),
                        open=float(row['open']),
                        high=float(row['high']),
                        low=float(row['low']),
                        close=float(row['close']),
                        vol=int(row['vol']),
                        amount=float(row['amount']) * 1000  # 千元转为元
                    )
                )

            # 批量插入（推荐使用execute批量操作）
            session.bulk_save_objects(records)
            session.commit()
            print(f"成功存储 {len(records)} 条记录")

        except Exception as e:
            session.rollback()
            print(f"数据存储失败: {str(e)}")
        finally:
            session.close()

    def save(self, data):
        """
            参数:
            data: 需要存储的数据
        """

        session = self.get_session()
        try:
            session.add(data)
            session.commit()
            print(f"成功存储 {data} 条记录")

        except Exception as e:
            session.rollback()
            print(f"数据存储失败: {str(e)}")
        finally:
            session.close()

    def bulk_save_objects(self, data: list):
        """
        批量插入 post 数据批量存储到MySQL

        参数:
            data: 需要存储的数据
        """
        if len(data) == 0:
            return

        session = self.get_session()
        try:
            session.bulk_save_objects(data)
            session.commit()
            print(f"成功存储 {len(data)} 条记录")

        except Exception as e:
            session.rollback()
            print(f"数据存储失败: {str(e)}")
        finally:
            session.close()

    def get_posts(self, page: int = 1, per_page: int = 20):
        """
        分页查询帖子数据

        参数:
            page (int): 当前页码
            per_page (int): 每页条数
        """
        session = self.get_session()
        try:
            offset = (page - 1) * per_page
            query = session.query(Post).order_by(Post.publish_time.desc())
            posts = query.offset(offset).limit(per_page).all()
            total = query.count()
            return {
                'items': posts,
                'page': page,
                'per_page': per_page,
                'total': total
            }
        finally:
            session.close()
